import mujoco_py
import numpy as np


def reset_mocap2body_xpos(simParam):
    """Resets the position and orientation of the mocap bodies to the same
    values as the bodies they're welded to.
    """
    if simParam.model.eq_type is None or simParam.model.eq_obj1id is None or simParam.model.eq_obj2id is None:
        return
    for eq_type, obj1_id, obj2_id in zip(simParam.model.eq_type, simParam.model.eq_obj1id, simParam.model.eq_obj2id):
        if eq_type != mujoco_py.const.EQ_WELD:
            continue
        mocap_id = simParam.model.body_mocapid[obj1_id]
        if mocap_id != -1:
            # obj1 is the mocap, obj2 is the welded body
            body_idx = obj2_id
        else:
            # obj2 is the mocap, obj1 is the welded body
            mocap_id = simParam.model.body_mocapid[obj2_id]
            body_idx = obj1_id
        assert mocap_id != -1
        simParam.data.mocap_pos[mocap_id][:] = simParam.data.body_xpos[body_idx]
        simParam.data.mocap_quat[mocap_id][:] = simParam.data.body_xquat[body_idx]


# load scene in MuJoCo
model = mujoco_py.load_model_from_path('assets/fetch/pick_and_place.xml')
sim = mujoco_py.MjSim(model, nsubsteps=20)
offscreen = mujoco_py.MjRenderContext(sim=sim, device_id=0, offscreen=True, opengl_backend='glfw')
viewer = mujoco_py.MjViewer(sim)
i = 0
while True:
    i += 1
    ctrl_action = np.array([1.0, 1.0])
    if sim.data.ctrl is not None:
        for i in range(ctrl_action.shape[0]):
            if sim.model.actuator_biastype[i] == 0:
                sim.data.ctrl[i] = ctrl_action[i]
            else:
                idx = sim.model.jnt_qposadr[sim.model.actuator_trnid[i, 0]]
                sim.data.ctrl[i] = sim.data.qpos[idx] + ctrl_action[i]
    mocap_action = np.array([0, 0, 1, 0, 0, 0, 0])
    if sim.model.nmocap > 0:
        mocap_action = mocap_action.reshape(sim.model.nmocap, 7)
        pos_delta = mocap_action[:, :3]
        quat_delta = mocap_action[:, 3:]
        reset_mocap2body_xpos(sim)
        sim.data.mocap_pos[:] = sim.data.mocap_pos + pos_delta
        sim.data.mocap_quat[:] = sim.data.mocap_quat + quat_delta
    sim.step()
    viewer.render()
